/**
 * 
 */
package com.gragra.sampling.nonproj;
import java.util.concurrent.atomic.AtomicInteger;
import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap;
import it.unimi.dsi.fastutil.ints.IntArrayList;
import it.unimi.dsi.fastutil.ints.IntIterator;
import it.unimi.dsi.fastutil.ints.IntList;
import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
import it.unimi.dsi.fastutil.ints.IntSet;
import com.gragra.sampling.CRPProcess;
import com.gragra.sampling.LogarithmicFactorialPlusAlphaCash;
import com.gragra.sampling.PopulationSamplingStructure;
import com.gragra.sampling.binary.HeadAssignmentStructure;
import com.gragra.sampling.nonproj.structures.nonprojective.NonProjDepdTree;
import com.gragra.sampling.sequentialStructure.WordAndTagHolder;
import com.gragra.util.OpenUnitElementCounter;
/**
 * @author Christoph Teichmann
 * 04.10.2013
 */
public class DMV extends CRPProcess implements ProposalDependencyAssigner
{
	/**
	 * 
	 */
	private int free = 0;
	/**
	 * 
	 * @return
	 */
	protected int getNextFree()
	{return free++;}
	/**
	 * 
	 */
	private final Int2IntOpenHashMap headToHeadLeftUnits;
	/**
	 * 
	 * @param head
	 * @param left
	 * @return
	 */
	protected int getHeadToHeadUnit(int head,boolean left)
	{
		if(left)
		{return headToHeadLeftUnits.get(head);}
		else
		{return headToHeadRightUnits.get(head);}
	}
	/**
	 * 
	 */
	private final Int2IntOpenHashMap headToHeadRightUnits;
	/**
	 * 
	 */
	private final Int2IntOpenHashMap stopAdjacentLeftUnits;
	/**
	 * 
	 */
	private final Int2IntOpenHashMap stopDistantLeftUnits;
	/**
	 * 
	 */
	private final Int2IntOpenHashMap stopAdjacentRightUnits;
	/**
	 * 
	 */
	private final Int2IntOpenHashMap stopDistantRightUnits;
	/**
	 * 
	 * @param head
	 * @param left
	 * @param adjacent
	 * @return
	 */
	protected int getStopUnit(int head, boolean left, boolean adjacent)
	{
		if(left)
		{
			if(adjacent)
			{return stopAdjacentLeftUnits.get(head);}
			else
			{return stopDistantLeftUnits.get(head);}
		}
		else
		{
			if(adjacent)
			{return stopAdjacentRightUnits.get(head);}
			else
			{return stopDistantRightUnits.get(head);}
		}
	}
	/**
	 * 
	 */
	private final IntSet sUnits;
	/**
	 * 
	 */
	private final IntSet hHUnits;
	/**
	 * 
	 */
	private final IntArrayList[] leftChildren;
	/**
	 * 
	 */
	private final IntArrayList[] rightChildren;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash hUnit;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash sUnit;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash hElement;
	/**
	 * 
	 */
	private final LogarithmicFactorialPlusAlphaCash sElement;
	/**
	 * @param population
	 */
	public DMV(int population,IntSet tags, double headHeadAlpha, double stopAlpha,double structMax, double structRed)
	{
		super(population);
		this.structuralStartingValue = structMax;
		this.structualReduction = structRed;
		this.sUnit = new LogarithmicFactorialPlusAlphaCash(stopAlpha*2, population);
		this.hUnit = new LogarithmicFactorialPlusAlphaCash(headHeadAlpha*tags.size(), population);
		this.hElement = new LogarithmicFactorialPlusAlphaCash(headHeadAlpha, population);
		this.sElement = new LogarithmicFactorialPlusAlphaCash(stopAlpha, population);
		IntIterator iit = tags.iterator();
		this.headToHeadLeftUnits = new Int2IntOpenHashMap();
		this.headToHeadRightUnits = new Int2IntOpenHashMap();
		this.stopAdjacentLeftUnits = new Int2IntOpenHashMap();
		this.stopAdjacentRightUnits = new Int2IntOpenHashMap();
		this.stopDistantLeftUnits = new Int2IntOpenHashMap();
		this.stopDistantRightUnits = new Int2IntOpenHashMap();
		while(iit.hasNext())
		{
			int tag = iit.nextInt();
			this.headToHeadLeftUnits.put(tag, this.getNextFree());
			this.headToHeadRightUnits.put(tag, getNextFree());
			this.stopAdjacentLeftUnits.put(tag, getNextFree());
			this.stopAdjacentRightUnits.put(tag, getNextFree());
			this.stopDistantLeftUnits.put(tag, getNextFree());
			this.stopDistantRightUnits.put(tag, getNextFree());
		}
		this.headToHeadLeftUnits.put(-1, this.getNextFree());
		this.headToHeadRightUnits.put(-1, getNextFree());
		this.stopAdjacentLeftUnits.put(-1, getNextFree());
		this.stopAdjacentRightUnits.put(-1, getNextFree());
		this.stopDistantLeftUnits.put(-1, getNextFree());
		this.stopDistantRightUnits.put(-1, getNextFree());
		this.leftChildren = new IntArrayList[population];
		this.rightChildren = new IntArrayList[population];
		this.seen = new OpenUnitElementCounter[population];
		for(int i=0;i<population;++i)
		{
			this.leftChildren[i] = new IntArrayList();
			this.rightChildren[i] = new IntArrayList();
			this.seen[i] = new OpenUnitElementCounter();
		}
		this.hHUnits = new IntOpenHashSet(this.headToHeadLeftUnits.values());
		this.hHUnits.addAll(this.headToHeadRightUnits.values());
		this.sUnits = new IntOpenHashSet(this.stopAdjacentLeftUnits.values());
		this.sUnits.addAll(this.stopAdjacentRightUnits.values());
		this.sUnits.addAll(this.stopDistantLeftUnits.values());
		this.sUnits.addAll(this.stopDistantRightUnits.values());
	}
	/**
	 * 
	 */
	private final OpenUnitElementCounter[] seen;
	/* (non-Javadoc)
	 * @see com.gragra.sampling.PopulationProbabilityAssigner#makeLogProb(com.gragra.sampling.PopulationSamplingStructure, int)
	 */
	@Override
	public double makeLogProb(PopulationSamplingStructure pss, int thread)
	{
		OpenUnitElementCounter oue = getSeen(thread);
		a.addAndGet(((HeadAssignmentStructure) pss).getLength());
		this.addFeatures(pss, oue, thread, 1.0);
		double d = this.makeLogProb(oue, thread);
		d += this.makeDistAndOtherConstraints(pss,thread);
		return d;
	}
	/**
	 * 
	 * @param pss
	 * @param thread
	 * @return
	 */
	protected double makeDistAndOtherConstraints(PopulationSamplingStructure pss,
			int thread)
	{
		/*
		double ret = 0.0;
		HeadAssignmentStructure has = (HeadAssignmentStructure) pss;
		for(int i=0;i<has.getLength();++i)
		{
			int head = has.getHeadPosition(i, thread);
			if(head != has.getDependencyRootPosition())
			{ret -= Math.pow(i-head, 2.0);}
		}*/
		return 0.0;
	}
	/**
	 * @param thread
	 * @return
	 */
	protected OpenUnitElementCounter getSeen(int thread)
	{
		OpenUnitElementCounter oue = seen[thread];
		oue.clear();
		return oue;
	}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.PopulationProbabilityAssigner#add(com.gragra.sampling.PopulationSamplingStructure, int)
	 */
	@Override
	public void add(PopulationSamplingStructure pss, int thread)
	{
		OpenUnitElementCounter oue = this.getGiven(thread);
		addFeatures(pss,oue,thread,1.0);
	}
	/**
	 * 
	 * @param pss
	 * @param oue
	 * @param thread
	 * @param dir
	 */
	protected void addFeatures(PopulationSamplingStructure pss,
			OpenUnitElementCounter oue, int thread, double dir)
	{
		makeLeftAndRightChildren(pss,thread);
		HeadAssignmentStructure has = (HeadAssignmentStructure) pss;
		WordAndTagHolder wth = (WordAndTagHolder) pss;
		for(int pos=0;pos<has.getLength();++pos)
		{
			int tag = wth.getTag(pos);
			int head = has.getHeadPosition(pos, thread);
			int headTag;
			if(head >= 0)
			{
				if(head == has.getDependencyRootPosition())
				{headTag = -1;}
				else
				{headTag = wth.getTag(head);}
				oue.add(this.getHeadToHeadUnit(headTag, pos < head), tag, dir);
			}
			int leftChildren = this.leftChildren[thread].getInt(pos);
			int rightChildren = this.rightChildren[thread].getInt(pos);
			makeStops(oue, dir, tag, leftChildren, true);
			makeStops(oue, dir, tag, rightChildren, false);
		}
	}
	/**
	 * @param oue
	 * @param dir
	 * @param tag
	 * @param leftChildren
	 * @param left
	 */
	protected void makeStops(OpenUnitElementCounter oue, double dir, int tag,
			int leftChildren, boolean left)
	{
		if(leftChildren == 0)
		{
			oue.add(this.getStopUnit(tag, left, true), 0, dir);
		}
		else
		{
			oue.add(this.getStopUnit(tag, left, true), 1, dir);
			if(leftChildren > 1)
			{oue.add(this.getStopUnit(tag, left, false), 1, dir*(leftChildren-1));}
			oue.add(this.getStopUnit(tag, left, false), 0, dir);
		}
	}
	/**
	 * 
	 * @param pss
	 * @param thread
	 */
	protected void makeLeftAndRightChildren(PopulationSamplingStructure pss,
			int thread)
	{
		IntArrayList left = this.leftChildren[thread];
		IntArrayList right = this.rightChildren[thread];
		left.clear();
		right.clear();
		HeadAssignmentStructure has = (HeadAssignmentStructure) pss;
		for(int pos=0;pos<has.getLength();++pos)
		{
			if(left.size()<=pos)
			{left.add(0);}
			if(right.size()<=pos)
			{right.add(0);}
			int head = has.getHeadPosition(pos, thread);
			if(head<0)
			{continue;}
			if(pos>head)
			{this.addChild(head,right);}
			else
			{this.addChild(head,left);}
		}
	}
	/**
	 * 
	 * @param head
	 * @param right
	 */
	private void addChild(int head, IntArrayList to)
	{
		while(to.size() <= head)
		{to.add(0);}
		to.set(head, to.get(head)+1);
	}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.PopulationProbabilityAssigner#remove(com.gragra.sampling.PopulationSamplingStructure, int)
	 */
	@Override
	public void remove(PopulationSamplingStructure pss, int thread)
	{
		OpenUnitElementCounter oue = this.getGiven(thread);
		addFeatures(pss,oue,thread,-1.0);
	}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.PopulationProbabilityAssigner#getLocalLogLike(com.gragra.sampling.PopulationSamplingStructure, int)
	 */
	@Override
	public double getLocalLogLike(PopulationSamplingStructure pss, int thread)
	{return this.makeDistAndOtherConstraints(pss, thread);}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.PopulationProbabilityAssigner#getNumberOfEvaluations()
	 */
	@Override
	public double getNumberOfEvaluations()
	{return a.get();}
	/**
	 * 
	 */
	private final AtomicInteger a = new AtomicInteger();
	/* (non-Javadoc)
	 * @see com.gragra.sampling.CRPProcess#getUnitCash(int)
	 */
	@Override
	protected LogarithmicFactorialPlusAlphaCash getUnitCash(int unit)
	{
		if(this.hHUnits.contains(unit))
		{return this.hUnit;}
		else if (this.sUnits.contains(unit)) 
		{return this.sUnit;}
		return null;
	}
	/* (non-Javadoc)
	 * @see com.gragra.sampling.CRPProcess#getElementCash(int, int)
	 */
	@Override
	protected LogarithmicFactorialPlusAlphaCash getElementCash(int unit,
			int element)
	{
		if(this.hHUnits.contains(unit))
		{return this.hElement;}
		else if (this.sUnits.contains(unit)) 
		{return this.sElement;}
		return null;
	}
	@Override
	public double makeProposalLogProb(int pos, int head,
			NonProjDepdTree tree, int thread)
	{
		this.a.addAndGet(-tree.getLength());
		this.auxs.addAndGet(tree.getLength());
		return this.makeLogProb(tree, thread);
	}
	/**
	 * 
	 */
	private final AtomicInteger auxs = new AtomicInteger();
	@Override
	public double makeChartProposalProb(
			NonProjDepdTree tree,
			int headCandidate, int child, int thread)
	{
		this.auxs.incrementAndGet();
		OpenUnitElementCounter oue = getSeen(thread);
		int tag = tree.getTag(child);
		int headTag;
		if(headCandidate == tree.getDependencyRootPosition())
		{headTag = -1;}
		else
		{headTag = tree.getTag(headCandidate);}
		oue.add(this.getHeadToHeadUnit(headTag, child < headCandidate), tag, 1.0);
		return this.makeLogProb(oue, thread);
	}
	@Override
	public double getNumberOfAuxEvals()
	{return this.auxs.get();}
	/**
	 * 
	 * @return
	 */
	protected IntList getChildren(boolean left, int thread)
	{
		if(left)
		{return this.leftChildren[thread];}
		else
		{return this.rightChildren[thread];}
	}
	/**
	 * 
	 */
	private double structAnnealing;
	/**
	 * 
	 */
	private final double structuralStartingValue;
	/**
	 * 
	 */
	private final double structualReduction;
	@Override
	public void reduceStructuralAnnealing(int p)
	{
		this.structAnnealing = this.structuralStartingValue-((p-1)*this.structualReduction);
		this.structAnnealing = this.structAnnealing < 0.0 ? 0.0 : this.structAnnealing;
	}
}